rm(list=ls())
library(data.table)
library(ggplot2)
library(polywog)
library(MASS)
library(doParallel)
library(directlabels)
load("ensemble.Rdata")
ensemble$sddistance <- ensemble$sd.distance/.12
ensemble <- ensemble[, sd.distance:=NULL]
ensemble <- ensemble[, diff:=parties-mean.num.gov.parties]
ensemble <- ensemble[,alpha5:=ifelse(alpha>.5, 1, 0)]

ensemble$sd.eccentricity <- ensemble$se.eccentricity*sqrt(500)
ensemble$sd.representation <- ensemble$se.representation*sqrt(500)
ensemble$sd.gov.representation2 <- ensemble$se.gov.representation2*sqrt(500)


## #########################
## # FUNCTION FOR PLOTTING #
## #########################
## plot.vars <- function (data, varying, by="",
##                          parties.val=5,
##                          discount.val=.5,
##                          alpha.val=.5,
##                          sd.distance.val=1
## ) {
##  if (varying=="parties"|by=="parties") {
##      parties.val <- unique(data[,parties])
##  }
##  if (varying=="discount"|by=="discount") {
##      discount.val <- unique(data[,discount])
##   }
##  if (varying=="sd.distance"|by=="sd.distance") {
##      sd.distance.val <- unique(data[,sd.distance])
##  }
##  if (varying=="rule"|by=="rule") {
##      rule.val <- unique(data[,rule])
##  }
##  if (varying=="alpha"|by=="alpha") {
##      alpha.val <- unique(data[,alpha])
##  }
##  output <- data[round==76&
##                    discount %in% discount.val
##                  & parties %in% parties.val
##                  & alpha %in% alpha.val
##                  & sd.distance %in% sd.distance.val]
##  output <- as.data.frame(output)
##  plot.data <-data.frame(x=output[,varying],
##                           y=output$mean,
##                           lower=output$mean-(1.96*output$se),
##                           upper=output$mean+(1.96*output$se),
##                           variable=output$variable,
##                           rule=output$rule)
##  if (by!="") {
##    plot.data$z <- output[,by]
##  } 
##      p <- ggplot(plot.data, aes(x=x, y=y, group=rule))
##      #p <- p + geom_segment(aes(colour=variable, y=upper, yend=lower, xend=x))
##      #p <- p + geom_point(aes(colour=variable))
##      p <- p + geom_ribbon(fill="grey73", aes(ymin=lower, ymax=upper))
##      p <- p + geom_line(aes(colour=as.factor(rule)))
##      p <- p + theme_bw()
##      p <- p + theme(panel.grid.major.x = element_blank())
##      p <- p + theme(panel.grid.major.y = element_blank())
##      p <- p + theme(panel.grid.minor.x = element_blank())
##      p <- p + theme(panel.grid.minor.y = element_blank())
##      p <- p + xlab(varying)
##      p <- p + ylab("")
##      p <- p + scale_colour_manual("Rules",
##                                   breaks=c(2, 3, 4, 7, 8),
##                                   values=c("red", "blue", "yellow", "green", "pink"),
##                                   labels=c("Aggregator", "Sat Governator", "Sticker",
##                                         "Hunter", "Hunt Governator"))
##      ## p <- p + scale_colour_manual("Rules", labels=c("Aggregator", "Sat Governator", "Sticker",
##      ##                                    "Hunter", "Hunt Governator"))
##      if (by!="") {
##          p <- p + facet_grid(variable~z, scales="free")
##       } else {
##          p <- p + facet_grid(variable~., scales="free")
##       }
## return(p)
## } #end plot.vars


## ## cl <- makeCluster(4)
## ## registerDoParallel(cl)
## ## formula=mean.eccentricity~alpha+alpha1+alpha0+
## ##                                           parties+
## ##                                           discount+discount0+discount1+
## ##                                           sddistance+sddistance0
## ## data=ensemble[rule==2&round==76]
## ## degree=3
## ## boot=1000
## ## parallel=TRUE
## ## select=.05
## ## lambda=NULL


## ## a <- polywog.select(mean.eccentricity~alpha+alpha1+alpha0+
## ##                                           parties+
## ##                                           discount+discount0+discount1+
## ##                                           sddistance+sddistance0, ensemble[rule==2&round==76], 3, 1000)


## #This function data-mines using polywog but eliminates terms from the equation
## polywog.select <- function(formula, data, degrees.cv, boot=5, select=.05, lambda=NULL) {
##     cv <- cv.polywog(formula=as.formula(formula),
##                      data=as.data.frame(data),
##                      degrees.cv=as.numeric(degrees.cv),
##                      .parallel=TRUE)
##     fit <- cv$polywog.fit
##     fit <- bootPolywog(fit, nboot = boot, .parallel=TRUE)  
##     sig <- which((abs(summary(fit)$coefficients[,1]/summary(fit)$coefficients[,2])>1.96)==TRUE)
##     rhs <- names(sig)
##     rhs <- gsub("[.]", ")*I(", rhs) 
##     if (rhs[1]=="(Intercept)") {
##         rhs <- rhs[-1]
##     } else {
##         rhs[1] <- "-1"
##     }
##     if (length(rhs)>1) {
##         rhs <- ifelse(rhs!="-1", paste("I(", rhs, ")", sep=""), "-1")
##         rhs <- paste(rhs, collapse="+")
##     } else if (length(rhs)==1) {
##         rhs <- paste("I(", rhs, ")", sep="")
##     } else if (length(rhs)==0){
##       rhs <- "1"
##     }  
##     new.formula <- as.formula(paste(as.character(formula[2]), "~", rhs))
##     return(lm(new.formula, data))
## }


predict.polywog <- function(pmodel, parties=3, discount=0, sddistance=0, alpha=0, memory=2, aspiration=.25, meanAshare=FALSE, 
                            meanSGshare=FALSE, meanHshare=FALSE, meanHGshare=FALSE) {
  #input.data <- expand.grid(0, 0, 0, parties, discount, sddistance, alpha, memory, aspiration, meanAshare, meanSGshare, meanHshare, meanHGshare, sharesticker)
  #names(input.data) <- c("eccentricity.mean", "representation.mean", "rep.gov.mean",
  if(length(parties)>1) {parties <- lapply(parties, function(x) expand.grid(x,seq(0,1,by=(1/x))))
                         parties <- do.call("rbind",parties)
                         #shares <- do.call("rbind",apply(parties, 1, functionX))
                         input.data <- expand.grid(0, parties[,1],discount, sddistance, alpha, memory, aspiration, 0,0,0,0)
                         input.data[,7 + which(c(meanAshare, meanSGshare, meanHshare, meanHGshare)==TRUE)] <- parties[,2]

  }
  if(length(parties)==1){  if(meanAshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, seq(0,1,1/parties), 0, 0, 0)}
                           if(meanSGshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, 0, seq(0,1,1/parties), 0, 0)}
                           if(meanHshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, 0, 0, seq(0,1,1/parties), 0)}
                           if(meanHGshare==TRUE){input.data <- expand.grid(0, parties,discount, sddistance, alpha, memory, aspiration, 0, 0, 0, seq(0,1,1/parties))}}
  names(input.data) <- c(names(pmodel$model)[1],                      
                         "parties", "discount", "sddistance", "alpha", "memory", "aspiration", "meanAshare",
                         "meanSGshare", "meanHshare", "meanHGshare")
  names(input.data)[1] <- ifelse(names(input.data)[1]=="log(mean.gov.representation2)", "mean.gov.representation2", names(input.data)[1])
  input.data <- as.data.frame(input.data)
  output.data <- as.data.frame(input.data)
  input.data$alpha1 <- ifelse(input.data$alpha==1, 1, 0)
  input.data$alpha0 <- ifelse(input.data$alpha==0, 1, 0)
  input.data$discount1 <- ifelse(input.data$discount==1, 1, 0)
  input.data$discount0 <- ifelse(input.data$discount==0, 1, 0)
  input.data$sddistance0 <- ifelse(input.data$sddistance==0, 1, 0)
  input.data$alpha5 <- ifelse(input.data$alpha>=.5,1,0)
  input.data$memory2 <- ifelse(input.data$memory==2,1,0)
  input.data$memory10 <- ifelse(input.data$memory==10,1,0)
  input.data$aspiration25 <- ifelse(input.data$aspiration==0.25,1,0)
  input.data$aspiration9 <- ifelse(input.data$aspiration==0.9,1,0)
  input <- model.matrix(pmodel, data=input.data)
  betas <- mvrnorm(1000, coef(pmodel), vcov(pmodel))
  values <- input %*% t(betas)
  means <- apply(values, 1, mean)
  lower <- apply(values, 1, quantile, p=.025)
  upper <- apply(values, 1, quantile, p=.975)
  return(as.data.table(cbind(output.data, lower, means, upper,
                             type=paste0(as.character(pmodel$terms[[2]]), collapse=""))))
}

library(ggplot2)
# Plotting predictions
plot.polywog <- function(x,parties, discount, sddistance, alpha, memory, aspiration, facets, plot=TRUE){
  pr.model2 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, TRUE, FALSE, FALSE, FALSE)
  pr.model3 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, FALSE, TRUE, FALSE, FALSE)
  pr.model4 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, FALSE, FALSE, TRUE, FALSE)
  pr.model5 <- predict.polywog(x,parties, discount, sddistance, alpha, memory, aspiration, FALSE, FALSE, FALSE, TRUE)
    means <- c(pr.model2$means, pr.model3$means, pr.model4$means, pr.model5$means)
  if(facets=="parties") { facet <- parties}
  if(facets=="discount") { facet <- discount}
  if(facets=="sddistance") { facet <- sddistance}
  if(facets=="alpha") { facet <- alpha}
  if(facets=="memory") { facet <- memory}
  if(facets=="aspiration") { facet <- aspiration}  
  shares <- rep(pr.model2$meanAshare,4)
  names <- c(rep("Aggregator",dim(pr.model2)[1]),rep("Satisficing\nGovernator",dim(pr.model2)[1]),rep("Hunter",dim(pr.model2)[1]),rep("Governator",dim(pr.model2)[1]))
  facet <-  rep(as.matrix(pr.model2[,facets, with=FALSE]),4)  
  upper <- c(pr.model2$upper, pr.model3$upper, pr.model4$upper, pr.model5$upper)
  lower <- c(pr.model2$lower, pr.model3$lower, pr.model4$lower, pr.model5$lower)
  data <- data.frame(means, names, shares, facet, upper, lower)
 if (plot==TRUE) {
  fig <- ggplot(data, aes(x=shares, y=means, group=names)) +
    facet_grid(~facet) +      
    geom_line(aes(colour=names))+
    geom_ribbon(aes(ymin=lower, ymax=upper), fill="grey73", alpha=0.2) +
    scale_x_continuous(breaks=c(0,0.5,1)) + 
    ylab("Predicted Difference to Sticker") + 
    theme_bw() +
    theme(legend.title=element_blank()) +
    theme(panel.grid.major.x = element_blank())+
    theme(panel.grid.major.y = element_blank())+
    theme(panel.grid.minor.x = element_blank())+
    theme(panel.grid.minor.y = element_blank())+
    theme(strip.background = element_blank())+
    scale_colour_manual("Rules",
                        values=c("red", "deeppink", "green", "blue"),
                        labels=c("Aggregator", "Governator", "Hunter",
                            "Satisficing\nGovernator"))
  return(fig)
 } else {
     return(data)
 }
}

load("analysis/sd_regression.model1_new.RData")
load("analysis/sd_regression.model2.Rdata")
load("analysis/sd_regression.model3.Rdata")

plot.data <- plot.polywog(model1_new, parties=5, discount=.5,
                          sddistance=1, alpha=0.5, memory=6, 
                     aspiration=.5, facets="parties", plot=FALSE)
plot.data <- as.data.table(plot.data)
plot.data$rule <- ifelse(plot.data$names=="Aggregator", "A", NA)
plot.data$rule <- ifelse(plot.data$names=="Governator", "G", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Hunter", "H", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Satisficing\nGovernator", "SG", plot.data$rule)
plot.data <- plot.data[,label.x:=1.03, list(names, facet)]
plot.data <- plot.data[,label.y:=means[shares==1], list(names, facet)]
plot.data$label.y <- ifelse(plot.data$rule=="A", plot.data$label.y-.05, plot.data$label.y)
plot.data$facet <- NULL
plot.data$facet <- "Variability in Government Misery:\nPredicted Difference to Sticker"
plot.data$plot <- 3
plot.data1 <- plot.data

plot.data <- plot.polywog(model2, parties=5, discount=.5, sddistance=1, alpha=0.5, memory=6, 
                     aspiration=.5, facets="parties", plot=FALSE)
plot.data <- as.data.table(plot.data)
plot.data$rule <- ifelse(plot.data$names=="Aggregator", "A", NA)
plot.data$rule <- ifelse(plot.data$names=="Governator", "G", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Hunter", "H", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Satisficing\nGovernator", "SG", plot.data$rule)
plot.data <- plot.data[,label.x:=1.03, list(names, facet)]
plot.data <- plot.data[,label.y:=means[shares==1], list(names, facet)]
plot.data$facet <- NULL
plot.data$facet <-"Variability in Party System Misery:\nPredicted Difference to Sticker"
plot.data$plot <- 2
plot.data2 <- plot.data



plot.data <- plot.polywog(model3, parties=5, discount=.5, sddistance=1, alpha=0.5, memory=6, 
                     aspiration=.5, facets="parties", plot=FALSE)
plot.data <- as.data.table(plot.data)
plot.data$rule <- ifelse(plot.data$names=="Aggregator", "A", NA)
plot.data$rule <- ifelse(plot.data$names=="Governator", "G", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Hunter", "H", plot.data$rule)
plot.data$rule <- ifelse(plot.data$names=="Satisficing\nGovernator", "SG", plot.data$rule)
plot.data <- plot.data[,label.x:=1.03, list(names, facet)]
plot.data <- plot.data[,label.y:=means[shares==1], list(names, facet)]
plot.data$label.y <- ifelse(plot.data$rule=="A", plot.data$label.y-.05, plot.data$label.y)
plot.data$facet <- NULL
plot.data$facet <-"Variability in Eccentricity:\nPredicted Difference to Sticker"
plot.data$plot <- 1
plot.data3 <- plot.data

plot.data <- as.data.table(rbind(plot.data1, plot.data2, plot.data3))
plot.data$plot <- factor(plot.data$plot,
                         labels=c("Variability in Eccentricity:\nPredicted Difference to Sticker",
                           "Variability in Party System Misery:\nPredicted Difference to Sticker",
                           "Variability in Government Misery:\nPredicted Difference to Sticker"))
## plot.data$plot <- relevel(plot.data$plot,
##                           "Variability in Party System Misery:\nPredicted Difference to Sticker")
## plot.data$plot <- relevel(plot.data$plot,
##                           "Variability in Eccentricity:\nPredicted Difference to Sticker")


fig3 <- ggplot(plot.data, aes(x=shares, y=means, group=names)) +
    geom_ribbon(aes(ymin=lower, ymax=upper), fill="grey73", alpha=0.2) +
    facet_grid(plot~., scales="free")+
    geom_line(aes(colour=names))+
    scale_x_continuous(breaks=c(0,0.5,1)) + 
    ylab("") + 
    theme_bw() +
    theme(legend.title=element_blank()) +
    theme(panel.grid.major.x = element_blank())+
    theme(panel.grid.major.y = element_blank())+
    theme(panel.grid.minor.x = element_blank())+
    theme(panel.grid.minor.y = element_blank())+
    theme(strip.background = element_blank())+
    theme(legend.position="none")+    
    geom_text(aes(x=label.x, y=label.y, label=rule))+
    theme(panel.margin = unit(.75, "lines"))+    
    xlab("Rule Shares of Non-Stickers")+    
    scale_colour_manual("Rules",
                        values=c("red", "deeppink", "green", "blue"),
                        labels=c("Aggregator", "Governator", "Hunter",
                            "Satisficing\nGovernator"))
fig3
ggsave("sd_evo.png", dpi=1000)

